{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "02e1009b",
   "metadata": {},
   "source": [
    "# mul 重写"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "69e59fb4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tvm\n",
    "from tvm import relax\n",
    "from tvm.relax.testing import nn\n",
    "from tvm_book.nn import modules\n",
    "from tvm.script import ir as I, relax as R, tir as T"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d1406d2",
   "metadata": {},
   "source": [
    "测试当乘法算子的右操作数不是常量时的量化处理逻辑。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "454b0f65",
   "metadata": {},
   "outputs": [],
   "source": [
    "bb = relax.BlockBuilder()\n",
    "\n",
    "with bb.function(\"main\"):\n",
    "    model = modules.Conv2D(16, 16, kernel_size=(3, 3), padding=(1, 1))\n",
    "    x = nn.Placeholder((1, 16, 64, 64), dtype=\"float32\", name=\"x\")\n",
    "    data = nn.Placeholder((1, 16, 1, 1), dtype=\"float32\", name=\"const\")\n",
    "    multiplier = R.sigmoid(data)\n",
    "    conv = model(x)\n",
    "    act = R.nn.relu(data=conv)\n",
    "    output = act * multiplier\n",
    "    params = [x, data] + model.parameters()\n",
    "    bb.emit_func_output(output, params)\n",
    "mod = bb.get()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "c2c90334",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import relax as R</span>\n",
       "\n",
       "<span style=\"color: #A2F\">@I</span><span style=\"color: #A2F; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #00F; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #A2F\">@R</span><span style=\"color: #A2F; font-weight: bold\">.</span>function\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">main</span>(x: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">64</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), const: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), bias: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">16</span>,)), weight: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>))) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">64</span>)):\n",
       "        gv: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">64</span>)) <span style=\"color: #A2F; font-weight: bold\">=</span> conv2_d(x, bias, weight)\n",
       "        gv1: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">64</span>)) <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>nn<span style=\"color: #A2F; font-weight: bold\">.</span>relu(gv)\n",
       "        gv2: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>sigmoid(const)\n",
       "        gv3: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">64</span>)) <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>multiply(gv1, gv2)\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> gv3\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "mod.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67c4a2e8",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py313",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
